package edu.jhu.nlp.sequence.classify.features;

/**
 * Adapted from Mallet's cc.mallet.share.casutton.ner.ConllNer2003Sentence2TokenSequence
 * This is bare bones version that just converts sentence to tokens 
 * TODO(delip): Remove unnecessary transforms
 * @author Delip Rao
 */



import cc.mallet.extract.StringSpan;
import cc.mallet.extract.StringTokenization;
import cc.mallet.pipe.Pipe;
import cc.mallet.types.*;

import java.util.regex.Pattern;

public class POSTokenizerPipe extends Pipe
{
  boolean saveSource = true;
  boolean doDigitCollapses = true;
  boolean doDowncasing = false;

  public POSTokenizerPipe()
  {
    super (null, new LabelAlphabet());
  }

  /* Lines look like this:
  -DOCSTART- -X- -X- O

  EU NNP
  rejects VBZ
  German JJ
  call NN
  to TO
  boycott VB
  British JJ
  lamb NN
  . . O

  ...
   */

  public Instance pipe (Instance carrier)
  {
    String sentenceLines = (String) carrier.getData();
    String[] tokens = sentenceLines.split ("\n");
    LabelSequence target = new LabelSequence ((LabelAlphabet)getTargetAlphabet(), tokens.length);
    StringBuffer source = saveSource ? new StringBuffer() : null;
    TokenSequence data = new StringTokenization (source);

    String word, label = null;

    for (int i = 0; i < tokens.length; i++) {
      if (tokens[i].length() != 0) {
        try {
          String[] features = tokens[i].split ("\\s");
          int fieldIdx = 0;
          word = features[fieldIdx++]; // .toLowerCase();
          if (features.length > 1) label = features[fieldIdx++];
        } catch (ArrayIndexOutOfBoundsException e) {
          throw new IllegalArgumentException ("Invalid line "+tokens[i]+" : expected word "
              + (isTargetProcessing () ? ", target" : "")
              + ".");
        }
      } else {
        word = "-<S>-";
        label = "O";
      }

      // Transformations
      if (doDigitCollapses) {
        if (word.matches ("19\\d\\d"))
          word = "<YEAR>";
        else if (word.matches ("19\\d\\ds"))
          word = "<YEARDECADE>";
        else if (word.matches ("19\\d\\d-\\d+"))
          word = "<YEARSPAN>";
        else if (word.matches ("\\d+\\\\/\\d"))
          word = "<FRACTION>";
        else if (word.matches ("\\d[\\d,\\.]*"))
          word = "<DIGITS>";
        else if (word.matches ("19\\d\\d-\\d\\d-\\d--d"))
          word = "<DATELINEDATE>";
        else if (word.matches ("19\\d\\d-\\d\\d-\\d\\d"))
          word = "<DATELINEDATE>";
        else if (word.matches (".*-led"))
          word = "<LED>";
        else if (word.matches (".*-sponsored"))
          word = "<LED>";
      }

      if (doDowncasing)
        word = word.toLowerCase();

      int start = source.length ();

      if (saveSource) {
        if (word.equals ("-<S>-")) source.append ("\n\n");
        source.append (word); source.append (" ");
      }

      Token token = new StringSpan (source, start, source.length () - 1);
      data.add (token);

      if (label != null) {
        target.add (label);
      }

    }

    carrier.setData(data);
    if (isTargetProcessing ()) carrier.setTarget(target);
    if (saveSource) carrier.setSource(source);
    return carrier;
  }

  private static final long serialVersionUID = 6748406892411164959L;

}